import numpy as np
import scipy


class KalmanFilter(object):
    """卡尔曼滤波器类"""
    def __init__(self):
        ndim, dt = 4, 1
        self.A = np.eye(2 * ndim, 2 * ndim)
        for i in range(ndim):
            self.A[i, ndim + i] = dt

        self.H = np.eye(ndim, 2 * ndim)

        self._std_weight_position = 1. / 20
        self._std_weight_velocity = 1. / 160

    def initiate(self, box: np.ndarray):
        """
        从未关联的 box 创建 track

        Parameters
            box:(x, y, a, h) with center position (x, y),
            aspect ratio a and height h

        Return
            (ndarray, ndarray): mean vector (8x1) and covariance matrix (8x8)
        """
        mean_pos = box
        mean_vel = np.zeros_like(mean_pos)
        mean = np.r_[mean_pos, mean_vel]

        std = [
            2 * self._std_weight_position * box[3],
            2 * self._std_weight_position * box[3],
            1e-2,
            2 * self._std_weight_position * box[3],
            10 * self._std_weight_velocity * box[3],
            10 * self._std_weight_velocity * box[3],
            1e-5,
            10 * self._std_weight_velocity * box[3]]
        covariance = np.diag(np.square(std))

        return mean, covariance

    def predict(self, mean, covariance):
        """
        卡尔曼滤波: 预测

        Parameters
            mean: ndarray(8x1)
            covariance: ndarray(8x8)

        Return
            (ndarray, ndarray): mean vector (8x1) and covariance matrix (8x8) of the predict step
        """
        std_pos = [
            self._std_weight_position * mean[3],
            self._std_weight_position * mean[3],
            1e-2,
            self._std_weight_position * mean[3]]
        std_vel = [
            self._std_weight_velocity * mean[3],
            self._std_weight_velocity * mean[3],
            1e-5,
            self._std_weight_velocity * mean[3]]
        Q = np.diag(np.square(np.r_[std_pos, std_vel]))

        mean = np.dot(self.A, mean)
        prior_estimation_covariance_k = np.linalg.multi_dot((
            self.A, covariance, self.A.T)) + Q

        return mean, prior_estimation_covariance_k

    def project(self, mean, covariance):
        """
        Parameters
            mean: ndarray(8x1)
            covariance: ndarray(8x8)

        Return
            (ndarray, ndarray): mean vector (8x1) and covariance matrix (8x8) of the predict step
        """
        std = [
            self._std_weight_position * mean[3],
            self._std_weight_position * mean[3],
            1e-1,
            self._std_weight_position * mean[3]]
        innovation_cov = np.diag(np.square(std))

        mean = np.dot(self.H, mean)
        covariance = np.linalg.multi_dot((
            self.H, covariance, self.H.T))

        return mean, covariance + innovation_cov

    def update(self, mean, covariance, box):
        """
        卡尔曼滤波: 更新

        Parameters
            mean: ndarray(8x1)
            covariance: ndarray(8x8)
            box: ndarray(4x1)
        Return
            (ndarray, ndarray):
        """
        projected_mean, projected_cov = self.project(mean, covariance)

        # chol_factor = np.linalg.cholesky(projected_cov)
        chol_factor, lower = scipy.linalg.cho_factor(
            projected_cov, lower=True, check_finite=False)
        kalman_gain = scipy.linalg.cho_solve(
            (chol_factor, lower), np.dot(covariance, self.H.T).T,
            check_finite=False).T
        innovation = box - projected_mean

        new_mean = mean + np.dot(innovation, kalman_gain.T)
        new_covariance = covariance - np.linalg.multi_dot((
            kalman_gain, projected_cov, kalman_gain.T))
        return new_mean, new_covariance

    def gating_distance(self, mean, covariance, measurements,
                        only_position=False):
        """Compute gating distance between state distribution and measurements.

        A suitable distance threshold can be obtained from `chi2inv95`. If
        `only_position` is False, the chi-square distribution has 4 degrees of
        freedom, otherwise 2.

        Parameters
        ----------
        mean : ndarray
            Mean vector over the state distribution (8 dimensional).
        covariance : ndarray
            Covariance of the state distribution (8x8 dimensional).
        measurements : ndarray
            An Nx4 dimensional matrix of N measurements, each in
            format (x, y, a, h) where (x, y) is the bounding box center
            position, a the aspect ratio, and h the height.
        only_position : Optional[bool]
            If True, distance computation is done with respect to the bounding
            box center position only.

        Returns
        -------
        ndarray
            Returns an array of length N, where the i-th element contains the
            squared Mahalanobis distance between (mean, covariance) and
            `measurements[i]`.

        """
        mean, covariance = self.project(mean, covariance)
        if only_position:
            mean, covariance = mean[:2], covariance[:2, :2]
            measurements = measurements[:, :2]

        cholesky_factor = np.linalg.cholesky(covariance)
        d = measurements - mean
        z = scipy.linalg.solve_triangular(
            cholesky_factor, d.T, lower=True, check_finite=False,
            overwrite_b=True)
        squared_maha = np.sum(z * z, axis=0)
        return squared_maha
